import torch
import torch.optim as opt
import numpy as np
import pandas as pd
import Attack_select.MySQLFunc_11 as mysql
import json
import Attack_select.Network as Network
import matplotlib.pyplot as plt
import sys

Data=json.load(open('Attack_select/Configs/Config_net.json','r',encoding='UTF-8'))
N_S,N_A,GAMMA,MAX_EP,Batch_Size,R_Span,Path=Data.values()

class Main():
    def __init__(self):
        Mysql=mysql.Mysql('主函数')
        conn=Mysql.get_conn()
        print('读取数据：',end='',flush=True)
        self.table=pd.read_sql('SELECT * FROM selectdatas.1v1_attack_case;',con=conn,index_col='ID')
        self.loss_record=[]
        print('完成,已加载',self.table.shape)

    def train(self):
        print('构建模型：',end='',flush=True)
        Net=Network.Net().cuda()
        Opt=opt.Adam(Net.parameters(),lr=0.0001)
        print('完成')
        for i in range(MAX_EP):
            data, label = self.Data_shample(Batch_Size)
            Opt.zero_grad()
            out=Net(data)
            loss=Net.loss(out,label)
            if i+1%R_Span==0:
                print(loss)
                self.loss_record.append(loss.tolist())
            loss.backward()
            Opt.step()
        torch.save(Net,Path)

    def Data_shample(self,Batch_size):
        Batch=self.table.sample(n=Batch_size)
        data=Batch.iloc[:,0:6]
        lable=Batch.iloc[:,6]
        lable=np.array(lable)
        #lable=np.eye(N_A)[lable.reshape(-1)]       # 转换onehot
        return np.array(data),np.array(lable)

    def render(self):
        x=range(0,MAX_EP,R_Span)
        plt.plot(x,self.loss_record)
        plt.show()

    def test(self):
        print('加载模型{}：'.format(Path),end='',flush=True)
        Net=torch.load(Path)
        print('完成')
        data, label = self.Data_shample(100)
        out = Net(data).tolist()
        out=self.out_process(out)
        plt.scatter(range(100),label,marker='*',c='red')
        plt.scatter(range(100), out, c='blue')
        plt.show()
        print('方案选择正确率：',self.list_diff(label,out))

    def test_correct_rate(self):
        print('加载模型{}：'.format(Path), end='',flush=True)
        Net = torch.load(Path)
        print('完成')
        record_correct_rate=[]
        print('Start:',end='',flush=True)
        for i in range(1000):
            data, label = self.Data_shample(100)
            out = Net(data).tolist()
            out = self.out_process(out)
            record_correct_rate.append(self.list_diff(label, out))
            if i%10==0:
                print('#',end='',flush=True)
        print('100%')
        rec=np.array(record_correct_rate)
        print('正确率均值：',rec.mean())
        print('正确率方差',rec.std())
        print('正确率最大：',rec.max())
        print('正确率最小：',rec.min())


    def out_process(self,x):
        out=[]
        x=np.array(x)
        for i in x:
            a=i.argmax()
            out.append(a)
        return out

    def list_diff(self,list1,list2):
        k=0
        for i,j in zip(list1,list2):
            if i==j or i==j+1 or i==j+2 or i==j-1 or i==j-2 or\
                i==j+3 or i==j-3 or i==j+4 or i==j-4 or i==j+5 \
                    or i==j-5 or i==j+6 or i==j-6:
                k+=1

        return k


if __name__=='__main__':
    M=Main()
    # M.train()
    # M.render()
    # M.test()
    M.test_correct_rate()